from exceptiongroup import catch
from typing_extensions import TypedDict
from server.reminder_maker.llms import LLMS
from langgraph.graph import END, StateGraph
import requests
from server.common.variable import THREAD_LOCAL_DATA

class GraphState(TypedDict):
    """
    Response the state of our graph.

    Attributes:
        task: reminder task
        cron_expression: cron expression for the reminder,
        reminder_text: reminder text
        res: result
    """
    task: str
    cron_expression: str
    reminder_text: str
    res: str


class ReminderMaker:

    def __init__(self, remake_threshold: int = 3):
        self.llms = LLMS()
        self.remake_threshold = remake_threshold

    def create_workflow(self):
        workflow = StateGraph(GraphState)
        # Define the node
        workflow.add_node("generate", self._generate)
        workflow.add_node("save_reminder", self._save_reminder)
        workflow.add_node("rewrite_task", self._rewrite_task)

        # Build graph
        workflow.set_entry_point("generate")
        workflow.add_conditional_edges("generate", self._decide_to_save,
                                       {
                                           "save reminder": "save_reminder",
                                           "rewrite task": "rewrite_task"
                                       })
        workflow.add_edge("rewrite_task", "generate")
        workflow.add_edge("save_reminder", END)

        return workflow.compile()

    # Node #
    def _generate(self, state):
        """
        Generate reminder
        Args:
            state (dict): The current graph state

        Returns:
            state (dict): New key added to state, corn_expression and reminder_text
        """
        print("---GENERATE REMINDER---")
        task = state["task"]

        # Get cron expression and reminder text from the LLM
        reminder = self.llms.reminder_generator.invoke({"task": task})

        return {"cron_expression": reminder.cron_expression, "reminder_text": reminder.reminder_text}

    def _save_reminder(self, state):
        """
        Request server to save reminder. After this, the reminder will run according to the cron expression.
        Args:
            state (dict): The current graph state

        Returns:
            state (dict): New key added to state, res
        """
        url = "http://127.0.0.1:9001/reminder/add"

        try:
            data = {"id": "", "username": "jim", "remindTime": state["cron_expression"], "content": state["reminder_text"]}
        except Exception as e:
            print(e)

        response = requests.post(url, data=data)

        if response.status_code == 200:
            print("---SAVE REMINDER SUCCESS---")
            res = response.text
        else:
            print("---SAVE REMINDER FAILURE---")
            res = "failed to save reminder, please try again."

        return {"res": res}

    def _rewrite_task(self, state):
        """
        rewrite reminder task
        Args:
            state (dict): The current graph state

        Returns:
            state (dict): new reminder task be added to state.
        """
        print("---REWRITE REMINDER TASK---")
        task = self.llms.task_rewriter.invoke({"task": state["task"]})

        return {"task": task.reminder_task}

    # Edge #
    def _decide_to_save(self, state):
        """
        Determines whether to save reminder, or re-generate a reminder task.

        Args:
            state (dict): The current graph state

        Returns:
            str: Binary decision for next node to call
        """
        score = self.llms.reminder_grader.invoke({"task": state["task"], "cron_expression": state["cron_expression"]})
        if score.binary_score >= self.remake_threshold:
            print("DECISION: SAVE REMINDER")
            return "save reminder"
        else:
            print("DECISION: REWRITE REMINDER TASK")
            return "rewrite task"